package geppetto.cat.programs;


import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;

import java.io.IOException;
import java.util.Random;



/**
 * class to compute alignment error statistics (AER, Precision, Recall & their
 * derivatives
 */
public class TresholdBySentenceLen {
	private String corpusDescription;

	private int size, maxSentenceSize;

	String modelDir;

	String modelName;

	String decodingType;

	float posterior;

	private AbstractModel mhmm;

	private BilingualCorpus forwardCorpus, backwardCorpus;

	public TresholdBySentenceLen(String[] args) {
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceSize = Integer.parseInt(args[2]);
		modelDir = args[3];
		modelName = args[4];
		decodingType = args[5];
		if (args.length > 6)
			posterior = Float.parseFloat(args[6]);
		else
			posterior = Float.NaN;
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("ModelDir " + modelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("Decoding type " + decodingType);
		System.out.println("---");
	}

	public void initializeCorpus(String corpusDesc, int size, int maxSentSize)
			throws IOException {
		System.out.println("Loading Corpus");
		forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDesc, size, maxSentSize);
		backwardCorpus = forwardCorpus.reverse();
		modelDir = modelDir + "/" + modelName + "/" + forwardCorpus.getName()
				+ "/" + size + "/model/";
		System.out.println("Finished Loading Corpus");
	}

	public void initializeModels(String modelName) {
		System.out.println("Loading Model");
		if (modelName.equalsIgnoreCase("baseline")) {
			mhmm = HMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
		} else if (modelName.equalsIgnoreCase("agreement")) {
			mhmm = AgreementHMM.loadModel(forwardCorpus, backwardCorpus,
					modelDir + "/MHMM/");
		} else {
			System.out.println("Unknown Model");
			System.exit(1);
		}
		System.out.println("Finished Loading Model");
	}

	public float tunePTS(int min, int max) {
		Random r = new Random(-9119629771948792409L);
		int[] sent = forwardCorpus.getSentencesByLen(
				BilingualCorpus.TRAIN_CORPUS, min, max);
		int totNumSent = sent.length;
		int[] sample = new int[200];
		if (totNumSent == 0)
			return 0; // Case where there are no sentences of that size
		for (int i = 0; i < sample.length; i++) {
			sample[i] = sent[r.nextInt(totNumSent)];
		}
		// count number of aligned points with viterbi
		int numAlignmentsF = mhmm.viterbiAlignments(sample,
				BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
		// tune posterior to get that.
		return mhmm.tuneThresholdNumPoints(numAlignmentsF, sample,
				BilingualCorpus.TRAIN_CORPUS,false);
	}

	public float tuneF1(int min, int max, float alpha) {
		int[] sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.DEV_CORPUS, min, max);
		System.out.println("Using to tune " + sentences.length);
		return mhmm.tuneTreshholdBalancedF1Aux(sentences,
				BilingualCorpus.DEV_CORPUS, alpha,false);
	}

	public float tuneAER(int min, int max) {
		int[] sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.DEV_CORPUS, min, max);
		System.out.println("Using to tune " + sentences.length);
		return mhmm.tuneTresholdAERAux(sentences, BilingualCorpus.DEV_CORPUS,false);
	}

	public void evaluateViter() {

		int[] sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.TEST_CORPUS, 0, 20);
		System.out.println("20 "
				+ AlignmentEvaluator.evaluate(
						mhmm.viterbiAlignments(sentences,
								BilingualCorpus.TEST_CORPUS),
						forwardCorpus.getAlignments(sentences,
								BilingualCorpus.TEST_CORPUS)).toSimpleString());
		sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.TEST_CORPUS, 20, 40);
		System.out.println("40 "
				+ AlignmentEvaluator.evaluate(
						mhmm.viterbiAlignments(sentences,
								BilingualCorpus.TEST_CORPUS),
						forwardCorpus.getAlignments(sentences,
								BilingualCorpus.TEST_CORPUS)).toSimpleString());
		sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.TEST_CORPUS, 40, Integer.MAX_VALUE);
		System.out.println("> 40 "
				+ AlignmentEvaluator.evaluate(
						mhmm.viterbiAlignments(sentences,
								BilingualCorpus.TEST_CORPUS),
						forwardCorpus.getAlignments(sentences,
								BilingualCorpus.TEST_CORPUS)).toSimpleString());
		sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.TEST_CORPUS, 0, Integer.MAX_VALUE);
		System.out.println("qll "
				+ AlignmentEvaluator.evaluate(
						mhmm.viterbiAlignments(sentences,
								BilingualCorpus.TEST_CORPUS),
						forwardCorpus.getAlignments(sentences,
								BilingualCorpus.TEST_CORPUS)).toSimpleString());
	}

	public Evaluation evalualte(int min, int max, float tresh) {
		int[] sentences = forwardCorpus.getSentencesByLen(
				BilingualCorpus.TEST_CORPUS, min, max);
		System.out.println("Using to evaluate " + sentences.length);
		AlignmentsSet al = mhmm.posteriorAlignments(sentences,
				BilingualCorpus.TEST_CORPUS, tresh,false,false);
		return AlignmentEvaluator.evaluate(al, forwardCorpus.getAlignments(
				sentences, BilingualCorpus.TEST_CORPUS));
	}

	public void evaluateSameTreshF1(float alpha) {
		float tresh = tuneF1(0, Integer.MAX_VALUE, alpha);
		System.out.println("20 " + evalualte(0, 20, tresh).toSimpleString());
		System.out.println("40 " + evalualte(20, 40, tresh).toSimpleString());
		System.out.println("> 40 "
				+ evalualte(40, Integer.MAX_VALUE, tresh).toSimpleString());
		System.out.println("all "
				+ evalualte(0, Integer.MAX_VALUE, tresh).toSimpleString());
	}

	public void evaluateDiffTreshF1(float alpha) {
		float tresh = tuneF1(0, 20, alpha);
		Evaluation total = new Evaluation(0, 0, 0, 0, 0, 0, 0);
		Evaluation eval1 = evalualte(0, 20, tresh);
		total.add(eval1);
		tresh = tuneF1(20, 40, alpha);
		Evaluation eval2 = evalualte(20, 40, tresh);
		total.add(eval2);
		tresh = tuneF1(40, Integer.MAX_VALUE, alpha);
		Evaluation eval3 = evalualte(40, Integer.MAX_VALUE, tresh);
		total.add(eval3);
		System.out.println("20 - " + eval1.toSimpleString());
		System.out.println("40 - " + eval2.toSimpleString());
		System.out.println("> 40 - " + eval3.toSimpleString());
		System.out.println("all " + total.toSimpleString());
	}

	public void evaluateSameTreshAER() {
		float tresh = tuneAER(0, Integer.MAX_VALUE);
		System.out.println("20 " + evalualte(0, 20, tresh).toSimpleString());
		System.out.println("40 " + evalualte(20, 40, tresh).toSimpleString());
		System.out.println("> 40 "
				+ evalualte(40, Integer.MAX_VALUE, tresh).toSimpleString());
		System.out.println("all "
				+ evalualte(0, Integer.MAX_VALUE, tresh).toSimpleString());
	}

	public void evaluateDiffTreshAER() {
		float tresh = tuneAER(0, 20);
		Evaluation total = new Evaluation(0, 0, 0, 0, 0, 0, 0);
		Evaluation eval1 = evalualte(0, 20, tresh);
		total.add(eval1);
		tresh = tuneAER(20, 40);
		Evaluation eval2 = evalualte(20, 40, tresh);
		total.add(eval2);
		tresh = tuneAER(40, Integer.MAX_VALUE);
		Evaluation eval3 = evalualte(40, Integer.MAX_VALUE, tresh);
		total.add(eval3);
		System.out.println("20 - " + eval1.toSimpleString());
		System.out.println("40 - " + eval2.toSimpleString());
		System.out.println("> 40 - " + eval3.toSimpleString());
		System.out.println("all " + total.toSimpleString());
	}

	public void evaluateSameTreshPTS() {
		float tresh = tunePTS(0, Integer.MAX_VALUE);
		System.out.println("20 " + evalualte(0, 20, tresh).toSimpleString());
		System.out.println("40 " + evalualte(20, 40, tresh).toSimpleString());
		System.out.println("> 40 "
				+ evalualte(40, Integer.MAX_VALUE, tresh).toSimpleString());
		System.out.println("all "
				+ evalualte(0, Integer.MAX_VALUE, tresh).toSimpleString());
	}

	public void evaluateDiffTreshPTS() {
		float tresh = tunePTS(0, 20);
		Evaluation total = new Evaluation(0, 0, 0, 0, 0, 0, 0);
		Evaluation eval1 = evalualte(0, 20, tresh);
		total.add(eval1);
		tresh = tunePTS(20, 40);
		Evaluation eval2 = evalualte(20, 40, tresh);
		total.add(eval2);
		tresh = tunePTS(40, Integer.MAX_VALUE);
		Evaluation eval3 = evalualte(40, Integer.MAX_VALUE, tresh);
		total.add(eval3);
		System.out.println("20 - " + eval1.toSimpleString());
		System.out.println("40 - " + eval2.toSimpleString());
		System.out.println("> 40 - " + eval3.toSimpleString());
		System.out.println("all " + total.toSimpleString());
	}

	public void run() throws IOException {
		initializeCorpus(corpusDescription, size, maxSentenceSize);
		initializeModels(modelName);

		System.out.println("viter\n\n");
		evaluateViter();

		System.out.println("AER\n\n");
		evaluateSameTreshAER();
		System.out.println("--------");
		evaluateDiffTreshAER();
		System.out.println("F1 0.3\n\n");
		evaluateSameTreshF1(0.3f);
		System.out.println("--------");
		evaluateDiffTreshF1(0.3f);
		System.out.println("F1 0.5\n\n");
		evaluateSameTreshF1(0.5f);
		System.out.println("--------");
		evaluateDiffTreshF1(0.5f);

		System.out.println("PTS\n\n");
		evaluateSameTreshPTS();
		System.out.println("--------");
		evaluateDiffTreshPTS();
	}

	public static void main(String[] args) throws IOException {
		TresholdBySentenceLen exp = new TresholdBySentenceLen(args);
		exp.run();
	}

}
